from PIL import Image, ImageDraw, ImageFont
import os
import platform

# --- 配置 ---
# 字体路径：自动适配多系统
SYSTEM = platform.system()
if SYSTEM == "Windows":
    FONT_PATH = "C:/Windows/Fonts/simhei.ttf"  # Windows黑体
elif SYSTEM == "Darwin":  # macOS
    FONT_PATH = "/Library/Fonts/Songti.ttc"  # macOS宋体
else:  # Linux
    FONT_PATH = "/usr/share/fonts/truetype/wqy/wqy-zenhei.ttc"

# 输入/输出目录
INPUT_DIR = "../add_text/pomeranian_colors"  # 原始图片目录（方形图片）
OUTPUT_DIR = "../add_text/pomeranian_colors_annotated_final"  # 输出目录
os.makedirs(OUTPUT_DIR, exist_ok=True)

# 毛色数据（与文件名对应）
color_data = [
    [1, "Orange", "橙色：颜色范围从深锈红色到浓郁的金色，再到浅蜂蜜金色。"],
    [2, "Cream", "奶油色：一种极其浅淡的橙色。"],
    [3, "Orange Sable", "橙色貂色：身体为橙色，底毛为更浅的橙色或奶油色，针毛尖端为黑色。"],
    [4, "Black and Tan", "黑褐色：基底为黑色，在特定部位（如眉毛、口吻、胸部、腿）有褐色斑块。"],
    [5, "Black", "黑色：底毛和针毛均为纯黑色。"],
    [6, "White", "白色：纯白色，不应有柠檬色或奶油色的阴影。"],
    [7, "Wolf Sable", "狼貂色：底毛为灰色，针毛为灰色带黑色尖端。"],
    [8, "Parti Color", "花色：白色为底色，带有其他颜色的斑块，头部有白色条纹为佳。"],
    [9, "Chocolate", "巧克力色：浓郁的棕色。"],
    [10, "Cream Sable", "奶油貂色：极浅的奶油色底毛，带有貂色（通常为黑色或深色）尖端。"],
    [11, "Blue Merle", "蓝色陨石色：黑色基底上带有灰色、蓝色斑块，像大理石花纹。"],
    [12, "Red", "红色：深锈红色。"],
    [13, "Chocolate and Tan", "巧克力褐：基底为巧克力色，带有褐色斑块。"],
    [14, "Red Sable", "红色貂色：深锈红色的底毛，带有貂色（黑色）尖端。"],
    [15, "Tri Colored", "三色：是花色博美，同时身上还带有褐色的斑块图案。"],
    [16, "Blue", "蓝色：是黑色的稀释色，呈现灰蓝色，鼻子、眼眶等应为蓝色。"],
    [17, "Blue Sable", "蓝色貂色：基底为蓝色，针毛为深蓝色，末端为黑色尖端。"],
    [18, "Blue and Tan", "蓝褐：基底为蓝色，带有褐色斑块。"],
    [19, "Chocolate Merle", "巧克力陨石色：巧克力色基底上带有浅棕色斑块，形成大理石花纹。"],
    [20, "Beaver", "海狸色/河狸色：颜色从浅饼干色到深棕灰色不等，毛发表面有银蓝色光泽。"],
    [21, "Black and Brindle", "黑虎斑：底色为金色、红色或橙色，带有强烈的黑色横条纹。"],
    [22, "Chocolate Sable", "巧克力貂色：基底为巧克力色，针毛为更深的巧克力色。"],
    [23, "Lavender", "薰衣草色：一种灰白色，毛发表面带有粉红-巧克力色调。"],
    [24, "Blue Brindle", "蓝虎斑：底色为金色、红色或橙色，带有强烈的蓝色横条纹。"],
    [25, "Merle and Tan", "陨石褐：在蓝色陨石色或巧克力陨石色的基础上，还带有褐色斑块。"]
]

# 文件名列表（与数据顺序对应）
actual_file_names = [
    "01_Orange.png", "02_Cream.png", "03_Orange_Sable.png", "04_Black_and_Tan.png", "05_Black.png",
    "06_White.png", "07_Wolf_Sable.png", "08_Parti_Color.png", "09_Chocolate.png", "10_Cream_Sable.png",
    "11_Blue_Merle.png", "12_Red.png", "13_Chocolate_and_Tan.png", "14_Red_Sable.png", "15_Tri_Colored.png",
    "16_Blue.png", "17_Blue_Sable.png", "18_Blue_and_Tan.png", "19_Chocolate_Merle.png", "20_Beaver.png",
    "21_Black_and_Brindle.png", "22_Chocolate_Sable.png", "23_Lavender.png", "24_Blue_Brindle.png",
    "25_Merle_and_Tan.png"
]

# 蒙版与文字样式配置
MASK_HEIGHT_RATIO = 0.2  # 蒙版高度占原图高度的20%
MASK_OPACITY = 180  # 蒙版透明度（0-255，180=半透明）
MASK_COLOR = (0, 0, 0)  # 蒙版颜色（黑色）
TEXT_COLOR = (255, 255, 255)  # 文字颜色（白色，与黑色蒙版对比）
FONT_SIZE_RATIO = 0.035  # 字体大小占原图宽度的比例
TEXT_PADDING_RATIO = 0.03  # 文字左右留白比例


# --- 工具函数 ---
def get_text_size(draw, text, font):
    """计算文本宽高"""
    bbox = draw.textbbox((0, 0), text, font=font)
    return (bbox[2] - bbox[0], bbox[3] - bbox[1])


def wrap_text(draw, text, font, max_width):
    """文本自动换行（适配中文）"""
    lines = []
    current_line = ""
    for char in text:
        test_line = current_line + char
        test_width = get_text_size(draw, test_line, font)[0]
        if test_width <= max_width:
            current_line = test_line
        else:
            lines.append(current_line)
            current_line = char
    if current_line:
        lines.append(current_line)
    return lines


def add_text_with_mask(input_path, output_path, text):
    """以蒙版方式在原图底部添加文字（不改变原图大小）"""
    try:
        with Image.open(input_path).convert("RGB") as img:
            width, height = img.size  # 保留原图尺寸
            print(f"处理图片：{os.path.basename(input_path)} | 尺寸：{width}x{height}")

            # 1. 创建蒙版层（与原图同尺寸，底部区域为半透明黑色）
            mask = Image.new('RGBA', (width, height), (0, 0, 0, 0))  # 全透明初始层
            mask_draw = ImageDraw.Draw(mask)

            # 计算蒙版区域（底部20%高度）
            mask_height = int(height * MASK_HEIGHT_RATIO)
            mask_y = height - mask_height  # 蒙版顶部Y坐标
            mask_draw.rectangle(
                [(0, mask_y), (width, height)],
                fill=(*MASK_COLOR, MASK_OPACITY)  # 半透明黑色蒙版
            )

            # 2. 准备文字（自动换行，最多2行）
            font_size = int(width * FONT_SIZE_RATIO)
            try:
                font = ImageFont.truetype(FONT_PATH, font_size)
            except IOError:
                font = ImageFont.load_default()
                print(f"警告：使用默认字体")

            # 计算文字最大宽度（减去左右留白）
            text_max_width = int(width * (1 - 2 * TEXT_PADDING_RATIO))
            draw_temp = ImageDraw.Draw(Image.new('RGB', (1, 1)))
            lines = wrap_text(draw_temp, text, font, text_max_width)

            # 限制最多2行，超出部分省略
            if len(lines) > 2:
                lines = lines[:2]
                lines[-1] += "..."  # 第二行末尾加省略号

            # 3. 计算文字位置（垂直居中于蒙版内）
            line_height = get_text_size(draw_temp, "测试", font)[1]
            line_spacing = int(line_height * 0.3)
            total_text_height = len(lines) * line_height + (len(lines) - 1) * line_spacing
            text_start_y = mask_y + (mask_height - total_text_height) // 2  # 蒙版内垂直居中
            text_x = int(width * TEXT_PADDING_RATIO)  # 左留白

            # 4. 在蒙版上绘制文字
            for i, line in enumerate(lines):
                text_y = text_start_y + i * (line_height + line_spacing)
                mask_draw.text((text_x, text_y), line, font=font, fill=TEXT_COLOR)

            # 5. 合并原图与蒙版（不改变原图大小）
            result = Image.alpha_composite(img.convert('RGBA'), mask)
            result = result.convert('RGB')  # 转成RGB格式保存（兼容JPG/PNG）
            result.save(output_path, quality=95)
            print(f"✅ 成功保存：{os.path.basename(output_path)}\n")

    except Exception as e:
        print(f"❌ 处理失败：{str(e)}\n")


# --- 批量处理 ---
if __name__ == "__main__":
    if not os.path.exists(INPUT_DIR):
        print(f"错误：输入目录 {INPUT_DIR} 不存在！")
        exit(1)

    # 校验数据与文件名数量一致
    if len(color_data) != len(actual_file_names):
        print(f"错误：数据({len(color_data)})与文件名({len(actual_file_names)})数量不匹配！")
        exit(1)

    # 遍历处理
    for idx, (data, filename) in enumerate(zip(color_data, actual_file_names)):
        color_num, color_en, color_cn = data
        input_path = os.path.join(INPUT_DIR, filename)
        output_path = os.path.join(OUTPUT_DIR, filename)

        if not os.path.exists(input_path):
            print(f"⚠️ 跳过：{filename} 未找到\n")
            continue

        # 构造文字内容（序号+名称+英文）
        text = f"{str(color_num).zfill(2)} {color_cn.split('：')[0]}（{color_en}）：{color_cn.split('：')[1]}"
        add_text_with_mask(input_path, output_path, text)

    print("=" * 50)
    print("处理完成！输出目录：", os.path.abspath(OUTPUT_DIR))
    print("=" * 50)